#!/usr/bin/env python3


import time
# import datetime
from bpfcc import BPF, USDT
import json
import atexit
# import platform
import subprocess
import sys
import os
from enum import IntEnum


class FunctionName(IntEnum):
    MUTEX_LOCK = 1
    MUTEX_UNLOCK = 2
    MUTEX_WRLOCK = 3
    MUTEX_RDLOCK = 4
    MUTEX_TIMEDLOCK = 5
    MUTEX_RWLOCK_UNLOCK = 6
    SEM_WAIT = 7
    SPIN_LOCK = 8

def get_function_name(func_type: int) -> str:
    """
    根据函数类型枚举值返回对应的函数名字符串。
    """
    try:
        func_enum = FunctionName(func_type)
        # 使用字典映射简化 if/elif 链 (可选优化)
        name_map = {
            FunctionName.MUTEX_LOCK: "pthread_mutex_lock",
            FunctionName.MUTEX_UNLOCK: "pthread_mutex_unlock",
            FunctionName.MUTEX_WRLOCK: "pthread_rwlock_wrlock",
            FunctionName.MUTEX_RDLOCK: "pthread_rwlock_rdlock",
            FunctionName.MUTEX_TIMEDLOCK: "pthread_mutex_timedlock", # 确认符号名
            FunctionName.MUTEX_RWLOCK_UNLOCK: "pthread_rwlock_unlock",
            FunctionName.SEM_WAIT: "sem_wait",
            FunctionName.SPIN_LOCK: "pthread_spin_lock",
        }
        return name_map.get(func_enum, "Unknown")
    except ValueError:
        return "Unknown"

def _convert_ktime_to_unix(bpf_ktime_ns):
    # 获取当前时间戳（CLOCK_REALTIME）和系统运行时间（uptime）
    current_realtime = time.time()
    try:
        with open("/proc/uptime", "r") as f:
            uptime_seconds = float(f.read().split()[0])
    except (IOError, ValueError) as e:
        print(f"Error reading /proc/uptime: {e}")
        return 0

    # 计算系统启动时间（CLOCK_REALTIME at boot）
    boot_time = current_realtime - uptime_seconds
    
    # 将 eBPF 时间戳转换为 Unix 时间戳（纳秒级）
    unix_ns = boot_time * 1e9 + bpf_ktime_ns
    return unix_ns

def _get_match_runtime_path(so_path, substr):
    for s in so_path:
        if substr in s:
            return s
    return None

def _get_host_runtime_path(directory, substr):
    """
    在指定目录及其子目录中递归查找 包括 substr的 .so 文件。

    :param directory: 要搜索的目录路径
    :return: 找到的第一个 libruntime.so 文件的绝对路径，如果没有找到则返回 None
    """
    so_path = []
    for root, dirs, files in os.walk(directory):
        if 'libruntime.so' in files:
            so_path.append(os.path.join(root, 'libruntime.so'))
    
    return _get_match_runtime_path(so_path, substr)

def _get_container_runtime_path(container_id, substr):
    try:
        # 获取容器的根文件系统路径
        inspect_command = f"docker inspect --format '{{{{.GraphDriver.Data.MergedDir}}}}' {container_id}"
        container_root = subprocess.check_output(inspect_command, shell=True).decode().strip()

        # 查找容器内的 libruntime.so文件
        find_command = f"find {container_root} -name libruntime.so"
        so_path = subprocess.check_output(find_command, shell=True).decode().strip()

        if so_path:
            return _get_match_runtime_path(so_path.split("\n"), substr)
        else:
            return "libruntime.so not found in the container."

    except subprocess.CalledProcessError as e:
        return f"Error: {e.output.decode().strip()}"

# Define the BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

#define MAX_ENTRIES 1024

enum lock_type_e {
    mutex_lock = 1,
    mutex_unlock,
    mutex_wrlock,
    mutex_rdlock,
    mutex_timedlock,
    mutex_rwlock_unlock,
    sem_wait,
    spin_lock
};

struct event {
    u64 start_ts;
    u64 end_ts;
    u32 pid;
    u32 tid;
    u64 lock;
    enum lock_type_e lock_type;
    char comm[TASK_COMM_LEN];
};

struct trace_lock {
    u64 tid;
    u64 start_ts;
    u64 lock_p;
};

BPF_HASH(start_times, u64, struct trace_lock);
BPF_PERF_OUTPUT(events);

int trace_lock_entry(struct pt_regs *ctx)
{
    struct trace_lock lock = {};

    u64 ts = bpf_ktime_get_ns();
    u64 tid = (u64)bpf_get_current_pid_tgid();

    lock.tid = tid;
    lock.start_ts = ts;
    lock.lock_p = (u64)PT_REGS_PARM1(ctx);

    // Store the timestamp and lock pointer in the hash map
    start_times.update(&tid, &lock);

    return 0;
}

static int trace_lock_exit_common(enum lock_type_e lock_type, struct pt_regs *ctx)
{
    struct event e = {};
    struct trace_lock *tlock = 0;
    u64 ts = 0;
    u64 tid = (u64)bpf_get_current_pid_tgid();

    // Get current timestamp in nanoseconds
    ts = bpf_ktime_get_ns();

    // Retrieve the start time from the hash map
    tlock = start_times.lookup(&tid);
    if (!tlock)
        return 0;   // Missed tracing context

    // Populate the event structure
    e.start_ts = tlock->start_ts;
    e.end_ts = ts;
    e.pid = tid >> 32;
    e.tid = (u32)tid;
    e.lock_type = lock_type;
    e.lock = tlock->lock_p;

    // Remove the entry from the hash map
    start_times.delete(&tid);

    // Output the event to user space
    DELAY_FILTER
    {
        bpf_get_current_comm(&e.comm, sizeof(e.comm));
        events.perf_submit(ctx, &e, sizeof(e));
    }
    return 0;
} 

int trace_mutex_lock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_lock, ctx);
}

int trace_mutex_unlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_unlock, ctx);
}

int trace_sem_wait_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(sem_wait, ctx);
}

int trace_spin_lock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(spin_lock, ctx);
}

int trace_mutex_wrlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_wrlock, ctx);
}

int trace_mutex_rdlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_rdlock, ctx);
}

int trace_mutex_timedlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_timedlock, ctx);
}

int trace_rwlock_unlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_rwlock_unlock, ctx);
}

"""

if len(sys.argv) != 4:
    print("Usage: python3 lock_trace.py <path_or_container_id> <lib_match_str> <delay_ms>")
    raise SystemExit(0)

path_or_container = sys.argv[1]
match_str = sys.argv[2]
delay_ms = 0
delay_ns = 0
try:
    delay_ms = int(sys.argv[3])
except ValueError:
    print("Error: delay_ms must be an integer.")
    raise SystemExit(1)

delay_ns = delay_ms * 1000 * 1000

if path_or_container[0] == '.' or path_or_container[0] == '/':
   lib_path = _get_host_runtime_path(path_or_container, match_str)
else:
   lib_path = _get_container_runtime_path(path_or_container, match_str)

if not lib_path:
    print("Error: Could not find the target .so file.")
    raise SystemExit(1)

print("Hook so: ", lib_path)

# Initialize BPF
b = BPF(text=bpf_text)

bpf_text = bpf_text.replace('DELAY_FILTER',
        'if ((e.end_ts - e.start_ts) > %d)' % delay_ns)

try:
    # Mutex
    b.attach_uprobe(name=lib_path, sym="pthread_mutex_lock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_mutex_lock", fn_name="trace_mutex_lock_exit")

    # Mutex unlock
    b.attach_uprobe(name=lib_path, sym="pthread_mutex_unlock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_mutex_unlock", fn_name="trace_mutex_unlock_exit")

    # Semaphore
    b.attach_uprobe(name=lib_path, sym="sem_wait", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="sem_wait", fn_name="trace_sem_wait_exit")

    # Spinlock
    b.attach_uprobe(name=lib_path, sym="pthread_spin_lock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_spin_lock", fn_name="trace_spin_lock_exit")

    # wrlock
    b.attach_uprobe(name=lib_path, sym="pthread_rwlock_wrlock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_wrlock", fn_name="trace_mutex_wrlock_exit")

    # rdlock
    b.attach_uprobe(name=lib_path, sym="pthread_rwlock_rdlock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_rdlock", fn_name="trace_mutex_rdlock_exit")

    # Timed lock
    b.attach_uprobe(name=lib_path, sym="pthread_mutex_timedlock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_mutex_timedlock", fn_name="trace_mutex_timedlock_exit")

    # rwlock unlock
    b.attach_uprobe(name=lib_path, sym="pthread_rwlock_unlock", fn_name="trace_lock_entry")
    b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_unlock", fn_name="trace_rwlock_unlock_exit")

except Exception as e:
    print(f"Error attaching uprobes/uretprobes: {e}")
    # 可以尝试列出库中的符号来调试
    # print(subprocess.check_output(["nm", "-D", lib_path]).decode())
    raise SystemExit(1)

# List to store events
events = []

def print_event(cpu, data, size):
    try:
        event = b["events"].event(data)
        decoded_func_name = get_function_name(event.lock_type)
        chrome_trace_event = {
            "name": decoded_func_name,
            "cat": "lock",
            "ph": "X",
            "pid": event.pid,
            "tid": event.tid,
            "ts": _convert_ktime_to_unix(event.start_ts) / 1000.0, # Convert ns to microseconds for Chrome Tracing
            "dur": (event.end_ts - event.start_ts) / 1000.0,       # Duration in microseconds
            "args": {
                "comm": event.comm.decode('utf-8', 'replace'),
                "lock_address": f"0x{event.lock:x}" # 以十六进制显示锁地址
            }
        }
        events.append(chrome_trace_event)
    except Exception as e:
        print(f"Error processing event: {e}")

# Register the callback function
b["events"].open_perf_buffer(print_event)

def dump_events():
    try:
        with open("log.json", "w") as f:
            json.dump(events, f, indent=2)
        print(f"\nEvents dumped to log.json (Total events: {len(events)})")
    except Exception as e:
        print(f"Error dumping events to file: {e}")

atexit.register(dump_events)

print(f"Tracing lock operations (delay > {delay_ms}ms) in '{lib_path}'... Hit Ctrl-C to end.")

# Loop to process events
while True:
    try:
        while True:
            b.perf_buffer_poll()
    except KeyboardInterrupt:
        print("\nDetaching probes and exiting...")
        # BCC 通常在对象销毁时会清理，但显式 detach 更好
        # 注意：BCC Python API 可能没有直接的 detach_all_uprobes 方法
        # 通常依赖于 BPF 对象销毁来清理
        # 如果需要更精细控制，可能需要手动跟踪和 detach
    finally:
        # 确保事件被转储
        dump_events()

# Usage (as comments):
# python3 lock_trace.py  <host_directory_path>  <library_substring> <min_delay_milliseconds>
# Example: python3 lock_trace.py  /usr/lib  "libruntime" 10
#
# python3 lock_trace.py  <container_id>  <library_substring> <min_delay_milliseconds>
# Example: python3 lock_trace.py  abc123def456  "libruntime" 50
